package org.loong.crypto.extension.cms;

import java.security.Key;
import java.util.HashMap;
import java.util.Map;

import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1Null;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.ASN1Primitive;
import org.bouncycastle.asn1.oiw.OIWObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.operator.GenericKey;
import org.loong.crypto.core.algorithm.EncryptionMethod;

public class EnvelopedDataHelper {

    protected static final Map<ASN1ObjectIdentifier, String> BASE_CIPHER_NAMES = new HashMap<>();

    protected static final Map<ASN1ObjectIdentifier, EncryptionMethod> CIPHER_ALG_NAMES = new HashMap<>();

    static {
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES128_ECB.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES192_ECB.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES256_ECB.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES128_CBC.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES192_CBC.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.AES256_CBC.getOID(), "AES");
        BASE_CIPHER_NAMES.put(EncryptionMethod.DESEDE_ECB.getOID(), "DESede");
        BASE_CIPHER_NAMES.put(EncryptionMethod.DESEDE_CBC.getOID(), "DESede");
        BASE_CIPHER_NAMES.put(EncryptionMethod.SM4_ECB.getOID(), "SM4");
        BASE_CIPHER_NAMES.put(EncryptionMethod.SM4_CBC.getOID(), "SM4");

        CIPHER_ALG_NAMES.put(EncryptionMethod.AES128_ECB.getOID(), EncryptionMethod.AES128_ECB);
        CIPHER_ALG_NAMES.put(EncryptionMethod.AES192_ECB.getOID(), EncryptionMethod.AES192_ECB);
        CIPHER_ALG_NAMES.put(EncryptionMethod.AES256_ECB.getOID(), EncryptionMethod.AES256_ECB);
        CIPHER_ALG_NAMES.put(EncryptionMethod.AES128_CBC.getOID(), EncryptionMethod.AES128_CBC);
        CIPHER_ALG_NAMES.put(EncryptionMethod.AES192_CBC.getOID(), EncryptionMethod.AES192_CBC);
        CIPHER_ALG_NAMES.put(EncryptionMethod.AES256_CBC.getOID(), EncryptionMethod.AES256_CBC);
        CIPHER_ALG_NAMES.put(EncryptionMethod.DESEDE_CBC.getOID(), EncryptionMethod.DESEDE_CBC);
        CIPHER_ALG_NAMES.put(EncryptionMethod.SM4_ECB.getOID(), EncryptionMethod.SM4_ECB);
        CIPHER_ALG_NAMES.put(EncryptionMethod.SM4_CBC.getOID(), EncryptionMethod.SM4_CBC);
    }

    public String getBaseCipherName(ASN1ObjectIdentifier algorithm) throws CMSException {
        if (!BASE_CIPHER_NAMES.containsKey(algorithm)) {
            throw new CMSException("Unsupported encryption algorithm: " + algorithm.getId());
        }

        return BASE_CIPHER_NAMES.get(algorithm);
    }

    public EncryptionMethod getEncryptionMethod(ASN1ObjectIdentifier algorithm) throws CMSException {
        if (!CIPHER_ALG_NAMES.containsKey(algorithm)) {
            throw new CMSException("Unsupported encryption algorithm: " + algorithm.getId());
        }

        return CIPHER_ALG_NAMES.get(algorithm);
    }

    public Key getKey(final ASN1ObjectIdentifier algorithm, GenericKey key) throws CMSException {
        if (key.getRepresentation() instanceof Key) {
            return (Key) key.getRepresentation();
        } else if (key.getRepresentation() instanceof byte[]) {
            return new SecretKeySpec((byte[]) key.getRepresentation(), getBaseCipherName(algorithm));
        } else {
            throw new IllegalArgumentException("unknown generic key type");
        }
    }

    public byte[] getIv(final Key encKey, final AlgorithmIdentifier encryptionAlgID) throws CMSException {
        ASN1Encodable parameters = encryptionAlgID.getParameters();
        if (parameters != null) {
            ASN1Primitive sParams = parameters.toASN1Primitive();
            if (sParams != null && !(sParams instanceof ASN1Null)) {
                return ASN1OctetString.getInstance(sParams).getOctets();
            } else {
                ASN1ObjectIdentifier encAlg = encryptionAlgID.getAlgorithm();
                if (encAlg.equals(OIWObjectIdentifiers.desCBC) || encAlg.equals(PKCSObjectIdentifiers.des_EDE3_CBC)) {
                    return new byte[8];
                } else {
                    return null;
                }
            }
        } else {
            return null;
        }
    }
}
